explanation_index = 0
model = bert_classifier
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
# We need to split forward pass into two part:
# 1) embeddings computation
# 2) classification
def compute_bert_outputs(model_bert, embedding_output, attention_mask=None, head_mask=None):
if attention_mask is None:
attention_mask = torch.ones(embedding_output.shape[0], embedding_output.shape[1]).to(embedding_output)
extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
extended_attention_mask = extended_attention_mask.to(dtype=next(model_bert.parameters()).dtype) # fp16 compatibility
extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
if head_mask is not None:
if head_mask.dim() == 1:
head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
head_mask = head_mask.expand(model_bert.config.num_hidden_layers, -1, -1, -1, -1)
elif head_mask.dim() == 2:
head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1) # We can specify head_mask for each layer
head_mask = head_mask.to(dtype=next(model_bert.parameters()).dtype) # switch to fload if need + fp16 compatibility
else:
head_mask = [None] * model_bert.config.num_hidden_layers
encoder_outputs = model_bert.encoder(embedding_output,
extended_attention_mask,
head_mask=head_mask)
sequence_output = encoder_outputs[0]
pooled_output = model_bert.pooler(sequence_output)
outputs = (sequence_output, pooled_output,) + encoder_outputs[1:] # add hidden_states and attentions if they are here
return outputs # sequence_output, pooled_output, (hidden_states), (attentions)
class BertModelWrapper(nn.Module):
def __init__(self, model):
super(BertModelWrapper, self).__init__()
self.model = model
def forward(self, embeddings):
outputs = compute_bert_outputs(self.model.bert, embeddings)
pooled_output = outputs[1] # Extract the pooled output tensor (usually at index 1)
# pooled_output = self.model.dropout(pooled_output) # Apply any required dropout
logits = self.model.classifier(pooled_output) # Pass the tensor to the classifier layer
# return torch.softmax(logits, dim=1)[:, 2].unsqueeze(1)
return torch.softmax(logits, dim=1)
# return outputs
bert_model_wrapper = BertModelWrapper(model)
ig = IntegratedGradients(bert_model_wrapper)
# accumalate couple samples in this array for visualization purposes
vis_data_records_ig = []
# Update the interpret_sentence function to move tensors to the same device as the model
def interpret_sentence(model_wrapper, sentence, label=1):
model_wrapper.eval()
model_wrapper.zero_grad()
# Tokenize and move input IDs to the device
input_ids = torch.tensor([tokenizer.encode(sentence, add_special_tokens=True, max_length=MAX_LEN, truncation=True)], device=device)
input_embedding = model_wrapper.model.bert.embeddings(input_ids)
# Ensure the model is on the same device
model_wrapper.to(device)
# Predict using the model wrapper
# pred = model_wrapper(input_embedding).detach().cpu().numpy()
# print("--->", pred.shape, pred[0][0])
# pred_ind = round(pred)
preds = model_wrapper(input_embedding).detach().cpu().numpy()
# print("shape", preds.shape, preds)
pred_ind = int(preds.argmax()) # Index of the predicted class
pred = preds[0][pred_ind]
# Compute attributions using Integrated Gradients
# attributions_ig, delta = ig.attribute(input_embedding, n_steps=500, return_convergence_delta=True)
attributions_ig, delta = ig.attribute(input_embedding, target=pred_ind, n_steps=100, return_convergence_delta=True)
print('pred: ', pred_ind, '(', '%.2f' % pred, ')', ', delta: ', abs(delta))
tokens = tokenizer.convert_ids_to_tokens(input_ids.cpu().numpy().flatten().tolist())
add_attributions_to_visualizer(attributions_ig, tokens, pred, pred_ind, label, delta, vis_data_records_ig)
def add_attributions_to_visualizer(attributions, tokens, pred, pred_ind, label, delta, vis_data_records):
attributions = attributions.sum(dim=2).squeeze(0)
attributions = attributions / torch.norm(attributions)
attributions = attributions.detach().cpu().numpy()
# storing couple samples in an array for visualization purposes
vis_data_records.append(visualization.VisualizationDataRecord(
attributions,
pred,
pred_ind,
label,
"label",
attributions.sum(),
tokens[:len(attributions)],
delta))
interpret_sentence(bert_model_wrapper, sentence=X_test[1], label=y_test[1])
_ = visualization.visualize_text(vis_data_records_ig)
